from bdtime import Time
from bdtime import tt
import os
import keyboard
from tools.mouse_state_recorder import DefaultArgs
import cv2
from tools import image_utils
from tools import get_time_id
from tools import get_folder_size
# from tools.enum_window import set_genshin_window_position
from tools import my_cv2_utils
from tools.remove_temp_file import remove_path


class ScreenRecorder:
    """
    录屏类
    """

    def __init__(self, save_file_name=None, operation_file_name=None, tt: (None, Time) = None,
                 resize=(600, 400), frames=DefaultArgs.frames,
                 wait_start_flag=False, check_stop_flag=False):
        self.tt: Time = tt if tt else Time()

        self.operation_file_name = operation_file_name
        self.wait_start_flag = wait_start_flag
        self.check_stop_flag = check_stop_flag

        self.frames = frames
        self.resize = resize

        self.time_id = get_time_id(self.tt)

        self.recorded_screen_images = []

        self.is_recording = False

        self.operation_save_dir = os.path.join(*['..', 'tempdir', 'recorded_states', 'screen_image', operation_file_name])
        os.makedirs(self.operation_save_dir, exist_ok=True)
        self._save_file_name = f'{len(os.listdir(self.operation_save_dir))}'
        self.save_file_name = save_file_name if save_file_name else self._save_file_name

        self.save_dir = os.path.join(*[self.operation_save_dir, self.save_file_name])
        print(f'====== ScreenRecorder.save_dir: {self.save_dir}')

        if os.path.exists(self.save_dir):
            remove_path(self.save_dir)
        os.makedirs(self.save_dir, exist_ok=True)

        self.stop_key = "alt + x"

    def get_save_file_path(self, file_name=None):
        return self.operation_save_dir
        # file_path = os.path.join(self.save_dir, file_name)
        # return file_path

    def stop_record(self):
        self.is_recording = False

    def record(self, until='alt + x'):
        keyboard.add_hotkey(until, self.stop_record)

        # tt = self.tt
        tt = Time()

        if self.wait_start_flag:
            while not DefaultArgs.flag__start_screen_record:
                tt.sleep(0.001)

        tt.__init__()
        self.is_recording = True
        i = 0
        print(f'~~~~~~ start screen_recorder.record, tt.now: {tt.now()}')

        while True:
            if self.check_stop_flag:
                if DefaultArgs.flag__stop_screen_record:
                    break

            if not self.is_recording or tt.stop(until, raise_error=False) or tt.stop('alt + x', raise_error=False):
                print(f'*** stop {self.__class__.__name__}.record by user!')
                break

            if DefaultArgs.flag__exit or DefaultArgs.flag__stop_screen_record:
                break

            next_record_time = DefaultArgs.record_time_schedule[i]
            while tt.now() < next_record_time:
                if DefaultArgs.use_sleep_wait:
                    tt.sleep(0.001)

            now = next_record_time
            # now = tt.now()

            img_arr = image_utils.cap_raw()

            if self.resize:
                img_arr = cv2.resize(img_arr, self.resize)

            img_name = f'screen__{i}__{now}.jpg'
            img_path = os.path.join(self.save_dir, img_name)

            my_cv2_utils.save_image(img_path, img_arr)

            self.recorded_screen_images.append(img_path)

            i += 1

        self.is_recording = False
        total_cost_time = tt.now()
        print(f'~~~~~~ end screen_recorder.record, tt.now: {total_cost_time}')

        # return self.recorded_screen_images
        return total_cost_time


def run(screen_recorder: ScreenRecorder = None, tt=None, save_file_name: str = None, resize=(600, 400), wait_start=False):
    if not screen_recorder:
        screen_recorder = ScreenRecorder(tt=tt, resize=resize, save_file_name=save_file_name)
    else:
        tt = screen_recorder.tt

    tt.sleep(1)
    print(f'====== start screen_recorder.record, tt.now: {tt.now()}')
    total_cost_time = screen_recorder.record()

    recorded_screen_size = get_folder_size(screen_recorder.save_dir)
    print(f'====== end screen_recorder.record --- size: {recorded_screen_size} --- total_cost_time: {total_cost_time} --- save_dir: {screen_recorder.save_dir}')


def main():
    # import keyboard
    # keyboard.press_and_release('win + d')

    # set_genshin_window_position(client_top_left=True, sticky=False)
    # set_genshin_window_position(client_top_left=False, sticky=False)

    tt = Time()
    resize = (600, 400)
    save_file_name = '蒙德传送点'
    # save_file_name = 'test'
    screen_recorder = ScreenRecorder(tt=tt, resize=resize, save_file_name=None, operation_file_name=save_file_name)
    run(screen_recorder=screen_recorder)


if __name__ == '__main__':
    main()
